package cn.guoshiweixue.security.utils;

import com.alibaba.fastjson2.JSONObject;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import java.security.*;
import java.security.spec.ECGenParameterSpec;
import java.util.Base64;

/**
 * RSA和SM2数字签名示例
 * @author guoshiweixue
 * @version 1.0
 * @date 2025/4/10
 */
public class DigitalSignatureWithJSON {
    static {
        // 注册 BouncyCastle 提供者
        Security.addProvider(new BouncyCastleProvider());
    }

    public static void main(String[] args) throws Exception {
        // 1. 生成RSA密钥对（公钥和私钥）
        KeyPair rsaKeyPair = generateKeyPair();
        PrivateKey rsaPrivateKey = rsaKeyPair.getPrivate(); // 私钥
        PublicKey rsaPublicKey = rsaKeyPair.getPublic();   // 公钥
        // 打印RSA公钥和私钥
        System.out.println("RSA Private Key: " + Base64.getEncoder().encodeToString(rsaPrivateKey.getEncoded()));
        System.out.println("RSA Public Key: " + Base64.getEncoder().encodeToString(rsaPublicKey.getEncoded()));

        // 2. 构造要签名的JSON对象
        JSONObject jsonObject = new JSONObject();
        jsonObject.put("id", 12345);
        jsonObject.put("name", "Alice");
        jsonObject.put("timestamp", System.currentTimeMillis());
        String originalMessage = jsonObject.toString();
        System.out.println("Original JSON Message: " + originalMessage);

        // 3. 使用RSA私钥对JSON对象进行签名
        String rsaSignature = signMessage(originalMessage, rsaPrivateKey);
        System.out.println("RSA Digital Signature: " + rsaSignature);

        // 4. 使用RSA公钥验证签名
        boolean rsaIsValid = verifySignature(originalMessage, rsaSignature, rsaPublicKey);
        System.out.println("Is the RSA signature valid? " + rsaIsValid);

        // 5. 生成SM2密钥对（公钥和私钥）
        KeyPair sm2KeyPair = generateSM2KeyPair();
        PrivateKey sm2PrivateKey = sm2KeyPair.getPrivate(); // 私钥
        PublicKey sm2PublicKey = sm2KeyPair.getPublic();   // 公钥

        // 打印SM2公钥和私钥
        System.out.println("SM2 Private Key: " + Base64.getEncoder().encodeToString(sm2PrivateKey.getEncoded()));
        System.out.println("SM2 Public Key: " + Base64.getEncoder().encodeToString(sm2PublicKey.getEncoded()));


        // 6. 使用SM2私钥对JSON对象进行签名
        String sm2Signature = signSM2Message(originalMessage, sm2PrivateKey);
        System.out.println("SM2 Digital Signature: " + sm2Signature);

        // 7. 使用SM2公钥验证签名
        boolean sm2IsValid = verifySM2Signature(originalMessage, sm2Signature, sm2PublicKey);
        System.out.println("Is the SM2 signature valid? " + sm2IsValid);
    }

    /**
     * 生成RSA密钥对（公钥和私钥）
     */
    public static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048); // 密钥长度为2048位
        return keyPairGenerator.generateKeyPair();
    }

    /**
     * 使用RSA私钥对消息进行签名
     */
    public static String signMessage(String message, PrivateKey privateKey) throws Exception {
        Signature signatureInstance = Signature.getInstance("SHA256withRSA");
        signatureInstance.initSign(privateKey);
        signatureInstance.update(message.getBytes());
        byte[] digitalSignature = signatureInstance.sign();
        return Base64.getEncoder().encodeToString(digitalSignature);
    }

    /**
     * 使用RSA公钥验证签名
     */
    public static boolean verifySignature(String message, String signature, PublicKey publicKey) throws Exception {
        Signature signatureInstance = Signature.getInstance("SHA256withRSA");
        signatureInstance.initVerify(publicKey);
        signatureInstance.update(message.getBytes());
        byte[] decodedSignature = Base64.getDecoder().decode(signature);
        return signatureInstance.verify(decodedSignature);
    }

    /**
     * 生成SM2密钥对（公钥和私钥）
     */
    public static KeyPair generateSM2KeyPair() throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, NoSuchProviderException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("EC", "BC"); // 或使用 "SM2"
        ECGenParameterSpec ecGenParameterSpec = new ECGenParameterSpec("sm2p256v1");
        keyPairGenerator.initialize(ecGenParameterSpec);
        return keyPairGenerator.generateKeyPair();
    }

    /**
     * 使用SM2私钥对消息进行签名
     */
    public static String signSM2Message(String message, PrivateKey privateKey) throws Exception {
        Signature signatureInstance = Signature.getInstance("SM3withSM2");
        signatureInstance.initSign(privateKey);
        signatureInstance.update(message.getBytes());
        byte[] digitalSignature = signatureInstance.sign();
        return Base64.getEncoder().encodeToString(digitalSignature);
    }

    /**
     * 使用SM2公钥验证签名
     */
    public static boolean verifySM2Signature(String message, String signature, PublicKey publicKey) throws Exception {
        Signature signatureInstance = Signature.getInstance("SM3withSM2");
        signatureInstance.initVerify(publicKey);
        signatureInstance.update(message.getBytes());
        byte[] decodedSignature = Base64.getDecoder().decode(signature);
        return signatureInstance.verify(decodedSignature);
    }
}
